package com.miplus.ccrm.core.api.aiyun;

import cn.hutool.json.JSONUtil;
import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.miplus.ccrm.core.api.aiyun.model.ChatRequest;
import lombok.extern.slf4j.Slf4j;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Sinks;

import java.util.Arrays;

@Slf4j
public class DashScopeApi {

    /**
     * 阿里云key
     */
    private static final String API_KEY = "sk-15cb2fe048c6409d9aa2ffb2ddf71bf6";


    /**
     * 阻塞式调用（非流式）
     */
    public static GenerationResult callWithMessage(ChatRequest chatRequest) throws ApiException, NoApiKeyException, InputRequiredException {
        Generation gen = new Generation();
        Message systemMsg = Message.builder()
                .role(Role.SYSTEM.getValue())
                .content(chatRequest.getSystemPrompt())
                .build();
        Message userMsg = Message.builder()
                .role(Role.USER.getValue())
                .content(chatRequest.getUserPrompt())
                .build();
        GenerationParam param = GenerationParam.builder()
                .apiKey(API_KEY)
                // 此处以qwen-plus为例，可按需更换模型名称。模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
                .model(chatRequest.getModel())
                .messages(Arrays.asList(systemMsg, userMsg))
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .build();
        GenerationResult result = gen.call(param);
        log.info("AI完整回答: {}", result.getOutput().getChoices().get(0).getMessage().getContent());
        return result;
    }

    /**
     * 流式调用
     */
    public static Flux<GenerationResult> streamCallWithFlux(String userMessage) {
        Sinks.Many<GenerationResult> sink = Sinks.many().replay().limit(1);
        StringBuilder fullContent = new StringBuilder(); // 用于存储拼接的完整内容

        new Thread(() -> {
            try {
                Generation gen = new Generation();
                Message systemMsg = Message.builder()
                        .role(Role.SYSTEM.getValue())
                        .content("You are a helpful assistant.")
                        .build();
                Message userMsg = Message.builder()
                        .role(Role.USER.getValue())
                        .content(userMessage)
                        .build();

                GenerationParam param = GenerationParam.builder()
                        .apiKey(API_KEY)
                        .model("qwen-plus")
                        .messages(Arrays.asList(systemMsg, userMsg))
                        .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                        .incrementalOutput(true)
                        .build();

                gen.streamCall(param).subscribe(
                        result -> {
                            // 打印每个流式片段（调试用）
                            log.info("Stream result: {}", JSONUtil.toJsonStr(result));
                            String incrementalContent = result.getOutput().getChoices().get(0).getMessage().getContent();
                            if (incrementalContent != null) {
                                fullContent.append(incrementalContent);
                            }
                            sink.tryEmitNext(result);
                        },
                        error -> {
                            log.error("Stream error: {}", error.getMessage());
                            sink.tryEmitError(error);
                        },
                        () -> {
                            log.info("Stream completed");
                            log.info("AI完整回答: {}", fullContent.toString());
                            sink.tryEmitComplete();
                        }
                );
            } catch (Exception e) {
                log.error("Exception in streaming thread: {}", e.getMessage());
                sink.tryEmitError(e);
            }
        }).start();
        return sink.asFlux();
    }
}
